# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import math
import torch.nn.functional as F
import torch
from torch import autograd

from fairseq import utils

from fairseq.criterions import FairseqCriterion, register_criterion
from .story_coherence import StoryCoherence
# from .story_cross_entropy_without_event_rl import StoryLabelSmoothedCrossEntropyCriterionWithoutEventRl
from .story_cross_entropy_without_event_rl_article import StoryLabelSmoothedCrossEntropyCriterionWithoutEventRlA


@register_criterion('story_rl')
class StoryRl(FairseqCriterion):

    def __init__(self, args, task):
        super().__init__(args, task)
        self.generator = StoryLabelSmoothedCrossEntropyCriterionWithoutEventRlA(args, task)
        self.discriminator = StoryCoherence(args, task)

    @staticmethod
    def add_args(parser):
        StoryLabelSmoothedCrossEntropyCriterionWithoutEventRlA.add_args(parser)

    @classmethod
    def build_criterion(cls, args, task):
        return cls(args, task)
